import os
import requests
import json
import io
import base64
from configparser import ConfigParser
from PIL import Image, PngImagePlugin
from gradio.processing_utils import encode_pil_to_base64
import datetime  

config = ConfigParser()
config.read("./src/arm.ini","UTF-8")


def time_stamp():
    # 获取当前日期和时间  
    now = datetime.datetime.now()  
    # 格式化日期和时间字符串  
    date_string = now.strftime("%Y-%m-%d-")  
    time_string = now.strftime("%H-%M-%S")  
    # 将日期和时间字符串组合在一起  
    combined_string = date_string + time_string  
    return combined_string


bat_path = './src/sd_init.bat'
url_txt2img = "http://localhost:7860/sdapi/v1/txt2img"
txt2img = {
            "enable_hr": True,
            "denoising_strength": 0.7,
            "firstphase_width": 0,
            "firstphase_height": 0,
            "prompt": "example prompt",
            "styles": [],
            "hr_second_pass_steps": 10,
            "hr_upscaler": "SwinIR_4x",
            "seed": -1,
            "subseed": -1,
            "subseed_strength": 0,
            "seed_resize_from_h": -1,
            "seed_resize_from_w": -1,
            "batch_size": 1,
            "n_iter": 1,
            "steps": config.getint('sd', 'steps'),
            "cfg_scale": 7,
            "width": config.getint('sd', 'width'),
            "height": config.getint('sd', 'height'),
            "restore_faces": bool(config['sd']['restore_faces']),
            "tiling": False,
            "negative_prompt": "",
            "eta": 0,
            "s_churn": 0,
            "s_tmax": 0,
            "s_tmin": 0,
            "s_noise": 1,
            "sampler_index": "DPM++ SDE Karras"
        }
requests.adapters.DEFAULT_RETRIES = 5
count = 0


def sd_init():
    os.system(bat_path)

def sd_work(positive_promt, negative_promt, url):
    global count
    count = count + 1
    txt2img["prompt"]=positive_promt+", <lora:child-drawing:1.3>,(no coloring:1.4), (white background:1.2), (smooth lines:1.1), easy to draw, black and white"
    txt2img["negative_prompt"]=negative_promt+",easynegative, (color:1.3), (nsfw:1.2), sexy,"
    #json_data = json.dump(txt2img)
    response = requests.post(url=f"{url}/sdapi/v1/txt2img", json=txt2img).json()
    for i in response['images']:
        image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
        png_payload = {
            "image": "data:image/png;base64,"+i
        }
        response2 = requests.post(url=f"{url}/sdapi/v1/png-info", json=png_payload)
        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("parameters", response2.json().get("info"))
        image.save(r'./tmp/pic_output/output'+time_stamp()+'.png', pnginfo=pnginfo)   
        image.save(r'./tmp/pic_output/output.png', pnginfo=pnginfo)   
    return r'./tmp/pic_output/output.png'
    

simple_img2img = {
            "init_images": None,
            "resize_mode": 0,
            "denoising_strength": 0.3,
            "mask": None,
            "mask_blur": 4,
            "inpainting_fill": 0,
            "inpaint_full_res": False,
            "inpaint_full_res_padding": 0,
            "inpainting_mask_invert": False,
            "prompt": "<lora:child-drawing:1>",
            "styles": [],
            "seed": -1,
            "subseed": -1,
            "subseed_strength": 0,
            "seed_resize_from_h": -1,
            "seed_resize_from_w": -1,
            "batch_size": 1,
            "n_iter": 1,
            "steps": 30,
            "cfg_scale": 3,
            "width": 512,
            "height": 512,
            "restore_faces": True,
            "tiling": False,
            "negative_prompt": "",
            "eta": 0,
            "s_churn": 0,
            "s_tmax": 0,
            "s_tmin": 0,
            "s_noise": 1,
            "override_settings": {},
            "sampler_index": "DPM++ SDE Karras",
            "include_init_images": False
            }

url_img2img = "http://localhost:7860/sdapi/v1/img2img"

def img2img(positive_promt, negative_promt, url,picpath):
    global count
    count = count + 1
    print(picpath)
    simple_img2img["init_images"]=[encode_pil_to_base64(Image.open(picpath))]
    simple_img2img["prompt"]=positive_promt+", <lora:child-drawing:1.3>,(no coloring:1.4), (white background:1.2), (smooth lines:1.1), easy to draw, black and white"
    simple_img2img["negative_prompt"]=negative_promt+",easynegative, (color:1.3), (nsfw:1.2), sexy,"
    #json_data = json.dump(txt2img)
    response = requests.post(url=f"{url}/sdapi/v1/img2img", json=simple_img2img).json()
    for i in response['images']:
        image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
        png_payload = {
            "image": "data:image/png;base64,"+i
        }
        response2 = requests.post(url=f"{url}/sdapi/v1/png-info", json=png_payload)
        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("parameters", response2.json().get("info"))
        image.save(r'./tmp/pic_output/output'+str(count)+'.png', pnginfo=pnginfo)   
        image.save(r'./tmp/pic_output/output.png', pnginfo=pnginfo)   
    return r'./tmp/pic_output/output.png'


def img2img_test():
    global count
    count = count + 1
    simple_img2img["prompt"]="<lora:child-drawing:1.3>"
    simple_img2img["negative_prompt"]="lowres, bad anatomy, text, error, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, artifacts, signature, watermark, username, blurry"
    #json_data = json.dump(txt2img)
    response = requests.post(url=f"http://localhost:7860/sdapi/v1/img2img", json=simple_img2img).json()
    for i in response['images']:
        image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
        png_payload = {
            "image": "data:image/png;base64,"+i
        }
        response2 = requests.post(url=f"http://localhost:7860/sdapi/v1/png-info", json=png_payload)
        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("parameters", response2.json().get("info"))
        image.save(r'./tmp/pic_output/output'+str(count)+'.png', pnginfo=pnginfo)
        image.save(r'./tmp/pic_output/output.png', pnginfo=pnginfo)
    return r'./tmp/pic_output/output.png'

if __name__ == "__main__":
    img2img_test()